import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense,Flatten
from tensorflow.keras.optimizers import RMSprop
import numpy as np
import matplotlib.pyplot as plt
import os
class Network(object):
    def __init__(self,train_x,train_y,test_x,test_y):
        self.train_x = train_x
        self.train_y = train_y
        self.test_x = test_x
        self.test_y = test_y
        self.model = self.build_network()
        self.labels = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
    'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
    def build_network(self):
        if os.path.exists('weight.h5'):
            model = keras.models.load_model('weight.h5')
        else:
            model = Sequential()
            model.add(Flatten(input_shape=(28,28)))
            model.add(Dense(128, activation='relu'))
            model.add(Dense(10, activation='softmax'))
            opt = RMSprop(lr=0.00001, epsilon=0.00001)
            model.compile(loss='sparse_categorical_crossentropy', optimizer=opt, metrics=['accuracy'])
            model.fit(self.train_x,self.train_y,batch_size=256,epochs=60)
            model.save('weight.h5')
        return model
    def train(self):
        loss,accuracy = self.model.evaluate(self.test_x,self.test_y)
        print('Train Loss:{},Accuracy:{}'.format(loss,accuracy))
        random_sample = 16
        rand_data = np.random.choice(1000,random_sample)
        count = 0

        test_sample = []
        test_label = []
        for i in rand_data:
            temp_x = self.test_x[i,:,:]
            test_sample.append(temp_x)

            temp_y = self.test_y[i]
            test_label.append(temp_y)
        test_sample = np.array(test_sample)
        result_predict = self.model.predict_classes(test_sample).tolist()

        fig_col = int(np.sqrt(random_sample))
        fig_row = random_sample // fig_col
        for i in range(random_sample):
            if test_label[i] == result_predict[i]:
                count += 1
            predict_label = self.labels[result_predict[i]]
            true_label = self.labels[test_label[i]]
            plt.subplot(fig_row,fig_col,i+1)
            plt.title(predict_label)
            print('Predict Label:{},True label:{}'.format(
                predict_label,true_label))
        print('Acuracy:{:.4f},Error nums:{}'.format(count/random_sample,random_sample-count))
        plt.savefig('../../figures/example/Predict.png',dpi=600)

def main():
    fashion_mnist = keras.datasets.fashion_mnist
    (train_x,train_y),(test_x,test_y) = fashion_mnist.load_data()
    train_x = train_x/255
    test_x = test_x/255
    # Show Image n
    nums = 16
    labels = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
    'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
    show_images = train_x[:nums,:,:]
    show_labels = train_y[:nums]
    for i in range(nums):
        plt.subplot(4,4,i+1)
        plt.imshow(show_images[i,:,:],cmap=plt.cm.binary)
        plt.title(labels[show_labels[i]])
    plt.savefig('../../figures/example/mnist_fashion.png',dpi=600)
    model = Network(train_x,train_y,test_x,test_y)
    model.train()
if __name__ == '__main__':
    main()
